import pandas as pd
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import adfuller
from statsmodels.tsa.arima.model import ARIMA

# 1. 读取CSV文件
data = pd.read_csv('data/数据挖掘/TSLA.csv')

# 2. 数据预处理
# 转换Date列为datetime类型，并设置为索引
data['Date'] = pd.to_datetime(data['Date'])
data = data.set_index('Date').asfreq('D')  # 设置日期频率为日，并填充缺失的日期

# 检查并处理缺失值，这里我们使用前一个有效值进行填充（前向填充）
data.fillna(method='ffill', inplace=True)

# 为了简化，我们选择'Close'列进行预测和可视化
close_prices = data['Close']

# 3. 使用ARIMA模型进行预测
# 首先检查数据的平稳性，如果不平稳则进行差分或其他转换
result = adfuller(close_prices)
print('ADF Statistic:', result[0])
print('p-value:', result[1])

# 假设数据已经平稳或经过差分后平稳，我们选择ARIMA模型进行预测
# 注意：在实际应用中，应该选择合适的ARIMA模型参数
# 这里为了简化，我们选择ARIMA(5, 1, 0)模型进行演示
model = ARIMA(close_prices, order=(15, 1, 0))
model_fit = model.fit()
forecast_steps = 15
forecast = model_fit.forecast(steps=forecast_steps)

# forecast 是一个 pandas Series，索引是日期，值是预测数据
# 使用 .values 属性获取预测值的一维数组
forecast_values = forecast.values

# 创建预测日期的索引，这里我们直接使用 forecast 的索引
forecast_index = forecast.index

# 数据可视化部分
plt.figure(figsize=(12, 6))
plt.plot(close_prices.index, close_prices.values, label='Actual Close Price', color='blue')
plt.plot(forecast_index, forecast_values, label='Forecasted Close Price', color='red', marker='o')
plt.legend()
plt.show()

